# coding=utf-8
import torch
import pandas as pd
import numpy as np
from tqdm import tqdm
"""
    这一部分代码用于实现基于item的协同过滤的相似度矩阵的生成
"""

# 余弦相似度（向量化）
def cosine_similarity_matrix(matrix_tensor):
    norm = matrix_tensor.norm(dim=1, keepdim=True)
    similarity_matrix = torch.mm(matrix_tensor, matrix_tensor.t()) / (norm * norm.t())
    return similarity_matrix
# 相似度矩阵
def similarity_matrix(matrix, save_path=None):
    matrix_tensor = torch.tensor(matrix.values, dtype=torch.float32)

    # 标准化每个用户向量为单位长度以改善余弦相似度计算
    matrix_tensor_normalized = matrix_tensor / matrix_tensor.norm(dim=1, keepdim=True)

    # 计算余弦相似度矩阵（使用向量化）
    similarity_matrix = cosine_similarity_matrix(matrix_tensor_normalized)

    similarity_df = pd.DataFrame(similarity_matrix.numpy(), index=matrix.index, columns=matrix.index)

    if save_path:
        similarity_df.to_csv(save_path)

    return similarity_df

# 读取数据集
train_data = pd.read_csv('./dataset/train_data.csv')

print("generate occurrence matrix:")
df = pd.DataFrame(train_data)
iu_co_matrix = pd.pivot_table(df, values='rating', index='book_id', columns='user_id', fill_value=0)

print("generate similarity matrix:")
similarity_matrix_save_path = 'E:/run_res/similarity_matrix_shuffle.csv'

# 生成并且保存矩阵
user_similarity_df = similarity_matrix(iu_co_matrix, save_path=similarity_matrix_save_path)

print("done!")